[dsv4] kv cache fp8#1600
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
ZhangLirong-amd
left a comment
ZhangLirong-amd
left a comment
There was a problem hiding this comment.
approve for running ci
9e09033 to
a4dbcb6
Compare
| """[padded_T+1] int32 GPU — per-token q indptr `arange(N+1)` (page_size=1, | ||
| max_seqlen_q=1). NOT `cu_seqlens_q` (which is per-seq and differs under | ||
| MTP); this is the per-token indptr the decode kernel consumes.""" | ||
| kv_last_page_lens: Optional[torch.Tensor] = None |
There was a problem hiding this comment.
asm decode kernel need them, so I builder them in metadata to aviod redundant opertaion at the layer level.
There was a problem hiding this comment.
I will update metadata once kernel update.
| self._rope_dtype = torch.bfloat16 # rope pool is always bf16 | ||
| # aiter prefill (op4) / decode (op5) support gfx950 (MI350) and | ||
| # gfx1250 (MI450) for the fp8 (2buff) path. | ||
| assert get_gfx() in ("gfx950", "gfx1250"), ( |
There was a problem hiding this comment.
if not support fallback to bf16, just same as what we have before this PR
| # region is self-consistent (qo_indptr[:N+1] == arange(N+1), kv_indptr[:N+1] | ||
| # is the real ragged cumsum) and the dropped tail is exactly the 0-length | ||
| # CG-padded slots. No-op when captured/padded (N == T_pad). | ||
| qo_indptr = qo_indptr[: N + 1].to(torch.int32) |
There was a problem hiding this comment.
remove these useless .to()
| # path leaves it in logits[:, 0]. logits.shape[1] = resolved split count. | ||
| resolved_splits = logits.shape[1] | ||
| if out_16_nosplit != 0 or resolved_splits > 1: | ||
| return output.to(torch.bfloat16) |
There was a problem hiding this comment.
This is an output guard mechanism. We can either adopt output guards or restrict input parameters of the interface.
Root cause: user set “out_16_nosplit=0 && resolved_splits==1”
| # and op5 (decode) consume directly — no torch quant, no in-kernel SWA | ||
| # scatter. | ||
| q_packed = q_rope_q = k_packed = k_rope = None | ||
| if self.kv_fp8: |
There was a problem hiding this comment.
merge these 2 path into 1
There was a problem hiding this comment.
Unify the interface for BF16 and FP8.
| self.attn_sink, | ||
| self.softmax_scale, | ||
| ) # [S, H, head_dim] | ||
| if self.kv_fp8: |
| # after this call and this avoids an extra empty_like allocation. | ||
| out=q_sa, | ||
| ) # [S, H, head_dim] | ||
| if self.kv_fp8: |
| swa_block_size, | ||
| min(self.window_size, attn_md.max_seqlen_q), | ||
| ) | ||
| if self.kv_fp8: |
4 participants
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist